This PR addresses issue #751 and #703

In #751 a user may want to compute complex comparisons or slopes that is not supported. Instead of adding more kwargs to handle each case, we add two functions data_grid, and select_draws that allow the user to compute complex quantities of interest.

data_grid allows the user to create their own pairwise grid (returned as a dataframe) which can then be passed to model.predict. Since the predict method returns the inference data, this simultaneously resolves #703.

Talking with @tomicapretto, we decide to not include a kwarg return_idata to comparisons, predictions, and slopes because if the user wants the inference data to compute their own quantities of interest, then they are only using the interpret module for its "data generating" capability, i.e., the data grid. Additionally, the user does not need to use data_grid, they can create their data however they deem fit, and pass it the predict method.

select_draws is a helper function that selects the posterior or posterior predictive draws from the InferenceData object returned by model.predict given a conditional dictionary. The conditional dictionary represents the values that correspond to that draw—this is essential in computing comparisons and slopes.

import warnings

import arviz as az
import numpy as np
import pandas as pd
import xarray as xr

import bambi as bmb

# the new functions
from bambi.interpret.helpers import data_grid, select_draws

bmb.config["INTERPRET_VERBOSE"] = False
warnings.simplefilter(action='ignore', category=FutureWarning)

fish_data = pd.read_stata("http://www.stata-press.com/data/r11/fish.dta")
cols = ["count", "livebait", "camper", "persons", "child"]
fish_data = fish_data[cols]
fish_data["child"] = fish_data["child"].astype(np.int8)
fish_data["persons"] = fish_data["persons"].astype(np.int8)
fish_data["livebait"] = pd.Categorical(fish_data["livebait"])
fish_data["camper"] = pd.Categorical(fish_data["camper"])

fish_model = bmb.Model(
    "count ~ livebait + camper + persons + child", 
    fish_data, 
    family='zero_inflated_poisson'
)

fish_idata = fish_model.fit(chains=4, random_seed=1234)

Below, we use the new data_grid to create a pairwise dataframe (cross-join) that will be the data for computing comparisons.

# create your own data
conditional = {
    "camper": np.array([0, 1]),
    "persons": np.arange(0, 5, 1),
}
variable = {"livebait": np.array([0, 1])}

# a pairwise grid
new_data = data_grid(fish_model, conditional, variable)
new_data.head()

Use this data in model.predict and then use the new select_draws function that selects the draws based on the conditional dictionary {"livebait": ...}. "count_mean" is the data variable we are selecting in the posterior group of the inference data. This returns the draws that satisfy the condition. Lastly, comparisons are computed.

idata_new = fish_model.predict(fish_idata, data=new_data, inplace=False)

draw_1 = select_draws(idata_new, new_data, {"livebait": 0}, "count_mean")
draw_2 = select_draws(idata_new, new_data, {"livebait": 1}, "count_mean")

# compute comparisons
(draw_2 - draw_1).mean(("chain", "draw"))

array([ 0.36152577,  0.86348116,  2.06643138,  4.95503272, 11.90508543,
        0.7086581 ,  1.69256968,  4.05050846,  9.71244674, 23.33495469])

This result above is the same as the estimate column below:

summary_df =bmb.interpret.comparisons(
    fish_model,
    fish_idata,
    contrast={"livebait": [0, 1]},
    conditional=conditional
)
summary_df

For a more advanced demo that computes cross-comparisons, see the advanced interpret usage docs (still a work in progress).

This is still a work in progress, so feedback is welcome. Don't dive too deep in a code review just yet. @tomicapretto @jt-lab @AylinH @zwelitunyiswa

To do:

• add tests
• tests need to pass
• resolve pylint errors
• resolve # TODO inline comments
• adapt interpret logger to work with the refactor
• add docs